#include <mosquittopp.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <fcntl.h>
#include <thread>
#include <mutex>
#include <atomic>
#include <iostream>
#include <cstring>
#include <functional>
#include <memory>
#include <map>

class MqttUdpBridge : public mosqpp::mosquittopp
{
public:
    MqttUdpBridge(const std::string& ipb, int portb, 
                  const std::string& ipu, int portuc, int portus,
                  const std::string& topic,
                  const std::map<short, std::string>& mPub, const std::map<std::string, short>& mSub)
        : mosquittopp("mqtt_udp_bridge"), 
          mqtt_broker_ip(ipb), mqtt_broker_port(portb),
          udp_target_ip(ipu), udp_target_port(portuc), 
          udp_listen_port(portus), 
          mSub(mSub), mPub(mPub),
          mqtt_sub_topic(topic),
          running(false)
    {
        
        // 初始化UDP目标地址
        memset(&udp_target_addr, 0, sizeof(udp_target_addr));
        udp_target_addr.sin_family = AF_INET;
        udp_target_addr.sin_port = htons(udp_target_port);
        inet_pton(AF_INET, udp_target_ip.c_str(), &udp_target_addr.sin_addr);
        
        // 初始化UDP监听地址
        memset(&udp_listen_addr, 0, sizeof(udp_listen_addr));
        udp_listen_addr.sin_family = AF_INET;
        udp_listen_addr.sin_port = htons(udp_listen_port);
        udp_listen_addr.sin_addr.s_addr = htonl(INADDR_ANY);
    }

    ~MqttUdpBridge()
    {
        stop();
    }

    bool start()
    {
        // 创建UDP socket
        udp_socket = ::socket(AF_INET, SOCK_DGRAM, 0);
        if (udp_socket < 0)
        {
            std::cerr << "Failed to create UDP socket" << std::endl;
            return false;
        }

        // 绑定UDP监听端口
        if (bind(udp_socket, (struct sockaddr*)&udp_listen_addr, sizeof(udp_listen_addr)) < 0)
        {
            std::cerr << "Failed to bind UDP socket to port " << udp_listen_port << std::endl;
            close(udp_socket);
            return false;
        }

        // 设置socket为非阻塞
        int flags = fcntl(udp_socket, F_GETFL, 0);
        fcntl(udp_socket, F_SETFL, flags | O_NONBLOCK);

        // 连接MQTT broker
        if (connect(mqtt_broker_ip.c_str(), mqtt_broker_port, 60) != MOSQ_ERR_SUCCESS)
        {
            std::cerr << "Failed to connect to MQTT broker" << std::endl;
            close(udp_socket);
            return false;
        }

        // 订阅主题
        if (subscribe(nullptr, mqtt_sub_topic.c_str()) != MOSQ_ERR_SUCCESS)
        {
            std::cerr << "Failed to subscribe to topic: " << mqtt_sub_topic << std::endl;
            close(udp_socket);
            return false;
        }

        running = true;
        
        // 启动UDP监听线程
        udp_listen_thread = std::thread(&MqttUdpBridge::udpListenLoop, this);
        
        // 启动MQTT循环（在调用线程中运行）
        loop_start();
        
        std::cout << "MQTT-UDP bridge started successfully" << std::endl;
        std::cout << "MQTT: " << mqtt_broker_ip << ":" << mqtt_broker_port << std::endl;
        std::cout << "Subscribe topic: " << mqtt_sub_topic << std::endl;
        std::cout << "UDP: listening on port " << udp_listen_port << ", sending to " << udp_target_ip << ":" << udp_target_port << std::endl;
        
        return true;
    }

    void stop()
    {
        if (running)
        {
            running = false;
            
            // 停止MQTT循环
            loop_stop();
            disconnect();
            
            // 停止UDP监听线程
            if (udp_listen_thread.joinable())
            {
                udp_listen_thread.join();
            }
            
            // 关闭socket
            if (udp_socket >= 0)
            {
                close(udp_socket);
            }
        }
    }

    // 获取当前状态信息
    std::string getStatus() const
    {
        return std::string("MQTT-UDP Bridge Status: ") + 
               (running ? "Running" : "Stopped") +
               "\nMQTT: " + mqtt_broker_ip + ":" + std::to_string(mqtt_broker_port) +
               "\nSubscribe: " + mqtt_sub_topic +
               "\nUDP: Listen " + std::to_string(udp_listen_port) + 
               " -> Target " + udp_target_ip + ":" + std::to_string(udp_target_port);
    }

private:
    // MQTT消息回调函数
    void on_message(const struct mosquitto_message* message) override
    {
        if (message->payloadlen > 0)
        {
            std::string topic(message->topic);
            auto it = mSub.find(topic);
            if (it != mSub.end())
            {
                std::string payload;
                payload.resize(message->payloadlen+2);
                char *p = (char*)payload.data();
                memcpy(p, &it->second, 2);
                memcpy(p+2, message->payload, message->payloadlen);
                sendUdpData(payload);
                static std::map<short,unsigned short>  mss;
                if(mss[it->second]++%100 == 0)
                {
                    printf("MQTT->UDP: %d\t%d\t%s\n", mss[it->second], message->payloadlen, it->first.c_str());
                }
            }
        }
    }

    void on_connect(int rc) override
    {
        if (rc == 0)
        {
            std::cout << "Connected to MQTT broker successfully" << std::endl;
        }
        else
        {
            std::cerr << "Failed to connect to MQTT broker, error code: " << rc << std::endl;
        }
    }

    void on_subscribe(int mid, int qos_count, const int* granted_qos) override
    {
        std::cout << "Subscribed to topic: " << mqtt_sub_topic << std::endl;
    }

    // 发送数据到UDP目标
    void sendUdpData(const std::string& data)
    {
        ssize_t sent = sendto(udp_socket, data.c_str(), data.length(), 0,
                             (struct sockaddr*)&udp_target_addr, sizeof(udp_target_addr));
        
        if (sent < 0)
        {
            std::cerr << "Failed to send UDP data" << std::endl;
        }
    }

    // UDP监听循环
    void udpListenLoop()
    {
        char buffer[4096];
        sockaddr_in client_addr;
        socklen_t client_len = sizeof(client_addr);

        while (running)
        {
            ssize_t received = recvfrom(udp_socket, buffer, sizeof(buffer) - 1, 0,
                                       (struct sockaddr*)&client_addr, &client_len);
            
            if (received > 0)
            {
                short id = *(short*)buffer;
                auto it = mPub.find(id);
                if(it != mPub.end())
                {
                    received -= 2;
                    static unsigned int sRecv = 0;
                    if (sRecv++%100 == 0)
                    {
                        std::cout << "UDP->MQTT: Received " << sRecv << ",\t" << received << "\tbytes, publishing to " << it->second << std::endl;
                    }
                    
                    // 发布到MQTT发布主题
                    int rc = publish(nullptr, it->second.c_str(), received, buffer+2, 0, false);
                    if (rc != MOSQ_ERR_SUCCESS)
                    {
                        std::cerr << "Failed to publish to MQTT, error code: " << rc << std::endl;
                    }
                }
            }
            else if (received < 0)
            {
                // 非阻塞socket，如果没有数据可读是正常的
                if (errno != EAGAIN && errno != EWOULDBLOCK)
                {
                    std::cerr << "UDP recvfrom error: " << strerror(errno) << std::endl;
                }
            }
            
            // 短暂休眠以避免CPU占用过高
            std::this_thread::sleep_for(std::chrono::milliseconds(10));
        }
    }

private:
    int udp_socket;
    int udp_target_port;
    int udp_listen_port;
    int mqtt_broker_port;
    struct sockaddr_in udp_target_addr;
    struct sockaddr_in udp_listen_addr;
    std::atomic<bool> running;
    std::thread udp_listen_thread;
    std::mutex mqtt_mutex;
    std::map<short, std::string> mPub;
    std::map<std::string, short> mSub;
    std::string mqtt_broker_ip;
    std::string mqtt_sub_topic;
    std::string udp_target_ip;
};

// 使用示例
int main(int argc, char **argv)
{
    char szMqttIP[16];
    char szWanIP[16];
    short nMqttPort = 1833;
    short nWanPortDst = 1122;
    short nWanPortRecv = 1123;

    std::map<short, std::string> mPub;
    std::map<std::string, short> mSub;
    FILE* fp = fopen("data.txt", "r");
    if (fp != NULL )
    {
        char line[256];       
        // 读取第一行
        fscanf(fp, "%s %hd %s %hd %hd", szMqttIP, &nMqttPort, szWanIP, &nWanPortDst, &nWanPortRecv);
        
        // 读取第二行及以后的行
        while (fgets(line, sizeof(line), fp))
        {
            short id;
            char szT1[64];
            char szT2[64];
            int count = sscanf(line, "%hd %s %s", &id, szT1, szT2);
        
            if (count >= 2)
            {
                mSub[szT1] = id;
                if (count >= 3)
                {
                    mPub[id] = szT2;
                }
            }
        }
        fclose(fp);
    }

    // 创建桥接实例 - 使用不同的订阅和发布topic
    MqttUdpBridge bridge(szMqttIP, nMqttPort,  // MQTT broker
        szWanIP, nWanPortDst,   // UDP目标
        nWanPortRecv,                    // UDP监听端口
        "bridge/+/+/+/+",
        mPub,        // 订阅的topic
        mSub);      // 发布的topic
    
    // 启动桥接
    if (!bridge.start())
    {
        std::cerr << "Failed to start bridge" << std::endl;
        return -1;
    }

    while (true)
    {
        std::this_thread::sleep_for(std::chrono::seconds(5));
    }
    
    return 0;
}

//echo -e -n "\x35\x82Connected to MQTT broker successfully" | socat - UDP4:127.0.0.1:1123
//Udp1122接收再发送给1123
//socat -d UDP-RECVFROM:1122,fork UDP4:127.0.0.1:1123